#ifndef MULTIGRIDONEDIM_H
#define MULTIGRIDONEDIM_H

#include "jacobi.hpp"
#include "option1D.hpp"
#include "norm.hpp"
#include <math.h>

using namespace std;
using namespace Eigen;

#define pi 4 * atan(1)
class Function
{
public:
    Function() {}
    double operator()(double _x)
    {
        return pow(_x, 2) - _x;
    }
    double f_x(double _x)
    {
        return 2 * _x - 1;
    }
    double laplace_f(double _x) // 负拉普拉斯算子
    {
        return -2;
    }
};
// class Function
// {
// public:
//     Function() {}
//     double operator()(double _x)
//     {
//         return 1.0 /(_x + 1.0);
//     }
//     double f_x(double _x)
//     {
//         return - 1.0 / pow(_x + 1.0, 2);
//     }
//     double laplace_f(double _x) // 负拉普拉斯算子
//     {
//         return - 2.0 /pow(_x + 1.0, 3);
//     }
// };

// class Function
// {
// public:
//     Function() {}
//     double operator()(double _x)
//     {
//         return exp(sin(_x));
//     }
//     double f_x(double _x)
//     {
//         return exp(sin(_x)) * cos(_x);
//     }
//     double laplace_f(double _x) // 负拉普拉斯算子
//     {
//         return exp(sin(_x)) * (sin(_x) - pow(cos(_x),2));
//     }
// };

/*****************************************一维多重网格*****************************************/
class MultigridOneDim
{
private:
    int n, dimention;
    double h, dh; // dh为全局h
    Function func;
    MatrixXd A, Ah;//A为全局
    VectorXd f, sol,//f,sol为全局
        vec_f, vec_sol;

public:
    /****************************************************************************Set Grid**********************************************************************************/
    void set_grid(BoundCondiType bound, int _n);

    /****************************************************************************V-cycle**********************************************************************************/
    VectorXd V_cycle(double w, VectorXd v, VectorXd _f, int nu1, int nu2,
                     BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter)
    {
        set_grid(bound, v.size() - 1);
        /***********************VC-1:Relax nu1 times***********************/
        iteration I(Ah, _f, w);
        v = I.weighted_jacobi_iteration(v, nu1);
        /***********************GridSize Judgement***********************/
        if ((v.size() - 1) != 2) // not Coarest Grid
        {
            int m = (v.size() - 1) / 2 + 1;
            VectorXd v2 = VectorXd::Zero(m);
            /***********************Restriction***********************/
            VectorXd r = _f - Ah * v;
            VectorXd tmp = RestrictionMatrix(rest, v.size() - 1) * r;
            v2 = V_cycle(w, v2, tmp, nu1, nu2, bound, rest, inter);
            /***********************Interpolation***********************/
            v = InterpolationMatrix(inter, m - 1) * v2 + v;
        }
        /***********************Relax nu1 times***********************/
        // is coarsest grid
        v = I.weighted_jacobi_iteration(v, nu2);
        return v;
    }
    /****************************************************************************FMG**********************************************************************************/
    VectorXd FMG(double w, VectorXd _f, int nu1, int nu2,
                 BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter)
    {
        VectorXd vtmp;
        if ((_f.size() - 1) != 2)
        {
            int m = (_f.size() - 1) / 2 + 1;
            VectorXd ftmp = RestrictionMatrix(rest, _f.size() - 1) * _f;
            vtmp = FMG(w, ftmp, nu1, nu2, bound, rest, inter);
            vtmp = InterpolationMatrix(inter, m - 1) * vtmp;
        }
        else
        {
            vtmp = VectorXd::Zero(_f.size());
        }
        vtmp = V_cycle(w, vtmp, _f, nu1, nu2, bound, rest, inter);
        return vtmp;
    }

    /****************************************************************************Multigrid**********************************************************************************/
    void Possion_BVP_Multigrid(int size, MultigridMethod method, double w, int nu1, int nu2,
                               BoundCondiType bound, Operators::Restriction rest, Operators::Interpolation inter,
                               int max_iter, double epsi)
    {//n, dh, f, sol, A
        set_grid(bound, size);
        n = size;
        dh = 1.0 / size;
        f = vec_f;
        sol = VectorXd::Zero(f.size());
        A = Ah;//A全局矩阵
        VectorXd Sol = A.lu().solve(vec_f);
        Norm Err;
        VectorXd real = VectorXd::Zero(f.size()); // 真实解
        double derr, Derr;
        for (int i = 0; i <= n; i++)
        {
            real(i) = func(i * dh);
        }

        if (method == MultigridMethod::Vcycle)
        {
            int cnt = 0;
            for (int i = 1; i <= max_iter; i++)
            {
                sol = V_cycle(w, sol, f, nu1, nu2, bound, rest, inter);
                derr = Err.inftynorm(sol - real);
                Derr = Err.inftynorm(Sol - sol);
                cnt++;
                cout << "V-cycle: " << i << "th: " << "解析解误差： " << derr << " 离散解误差： "<< Derr << endl;
                //cout << real;
                if (Derr < epsi)
                    break;
            }
            cout << "err < epsi after " << cnt << " iteration" << endl;
        }
        else if (method == MultigridMethod::FMG)
        {
            sol = FMG(w, f, nu1, nu2, bound, rest, inter);
            derr = Err.inftynorm(sol - real);
            Derr = Err.inftynorm(Sol - sol);
            cout << "FMG: " << 1 << "th:" << "解析解误差： " << derr << " 离散解误差： "<< Derr << endl;
        }
    }

    void display()
    {
        cout << "x = [";
        for (int i = 0; i <= n; i++)
        {
            cout << i * dh << ",";
        }
        cout << "];" << endl;

        cout << "y = [";
        for (int i = 0; i <= n; i++)
        {
            cout << sol[i] << ",";
        }
        cout << "];" << endl;

        cout << "scatter(x,y, \".\");" << endl;
    }
};

void MultigridOneDim::set_grid(BoundCondiType bound, int _n)
{
    h = 1.0 / _n;
    dimention = _n + 1;
    /***********************Dirichlet Boundary***********************/
    if (bound == BoundCondiType::Dirichlet)
    {
        Ah = MatrixXd::Zero(dimention, dimention);
        vec_f = VectorXd::Zero(dimention);
        // cout << A << endl;
        for (int i = 0; i < dimention; i++)
        {
            if ((i == 0) || (i == dimention - 1))
            {
                Ah(i, i) = 1;
                vec_f(i) = func(i * h);
            }
            else
            {
                Ah(i, i - 1) = -1.0 / pow(h, 2);
                Ah(i, i + 1) = -1.0 / pow(h, 2);
                Ah(i, i) = 2.0 / pow(h, 2);
                vec_f(i) = func.laplace_f(i * h);
            }
        }
        //cout << vec_f << endl;
    }
    /***********************Neumann Boundary***********************/
    else if (bound == BoundCondiType::Neumann)
    {
        Ah = MatrixXd::Zero(dimention, dimention);
        vec_f = VectorXd::Zero(dimention);
        for (int i = 0; i < dimention; i++)
        {
            if (i == 0)
            {
                Ah(i, i) = 1;
                vec_f(i) = func(0);
            }
            else if (i == dimention - 1)
            {
                Ah(i, i) = 1.0 / h;
                Ah(i, i - 1) = -1.0 / h;
                vec_f(i) = func.f_x(1.0);
            }
            else
            {
                Ah(i, i) = 2.0 / pow(h, 2);
                Ah(i, i - 1) = -1.0 / pow(h, 2);
                Ah(i, i + 1) = -1.0 / pow(h, 2);
                vec_f(i) = func.laplace_f(i * h);
            }
        }
    }
}

#endif